#!/usr/bin/env python3

import rospy
import yaml
import os
from visualization_msgs.msg import MarkerArray, Marker
from geometry_msgs.msg import Point

def load_graph_data(yaml_path):
    """加载YAML格式的节点地图数据"""
    if not os.path.exists(yaml_path):
        rospy.logerr(f"YAML file is not exist: {yaml_path}")
        return None, None
    
    with open(yaml_path, 'r') as f:
        try:
            graph_data = yaml.safe_load(f)  # 使用safe_load提高安全性
            return graph_data['nodes'], graph_data['edges']
        except yaml.YAMLError as e:
            rospy.logerr(f"fail to extract YAML: {e}")
            return None, None

def create_marker_array(nodes, edges, frame_id):
    """创建可视化MarkerArray"""
    marker_array = MarkerArray()

   # 节点标记（红色点）
    nodes_marker = Marker()
    nodes_marker.header.frame_id = frame_id
    nodes_marker.type = Marker.POINTS
    nodes_marker.action = Marker.ADD
    nodes_marker.pose.orientation.w = 1.0
    nodes_marker.scale.x = 0.2  # 点的宽度
    nodes_marker.scale.y = 0.2  # 点的高度
    nodes_marker.color.r = 1.0  # 红色
    nodes_marker.color.a = 1.0  # 不透明度
    nodes_marker.id = 0

    for node in nodes:
        point = Point()
        point.x, point.y, point.z = node[0], node[1], node[2]
        nodes_marker.points.append(point)
    marker_array.markers.append(nodes_marker)

    # 边标记（蓝色线）
    edges_marker = Marker()
    edges_marker.header.frame_id = frame_id
    edges_marker.type = Marker.LINE_LIST
    edges_marker.action = Marker.ADD
    edges_marker.pose.orientation.w = 1.0
    edges_marker.scale.x = 0.05  # 线宽
    edges_marker.color.b = 1.0  # 蓝色
    edges_marker.color.a = 1.0  # 不透明度
    edges_marker.id = 1

    for edge in edges:
        n1, n2 = edge
        if n1 < len(nodes) and n2 < len(nodes):  # 确保索引有效
            p1 = Point(x=nodes[n1][0], y=nodes[n1][1], z=nodes[n1][2])
            p2 = Point(x=nodes[n2][0], y=nodes[n2][1], z=nodes[n2][2])
            edges_marker.points.extend([p1, p2])
    marker_array.markers.append(edges_marker)

    return marker_array

def main():
    rospy.init_node('node_graph_loader', anonymous=True)

    # 获取参数
    map_path = rospy.get_param('~map_path')
    frame_id = rospy.get_param('~frame_id', 'map')
    topic_name = rospy.get_param('~topic_name', '/node_graph')
    rate_val = rospy.get_param('~rate', 1)

    # 构建YAML路径
    base, _ = os.path.splitext(map_path)
    yaml_path = base + '.yaml'

    # 加载数据
    nodes, edges = load_graph_data(yaml_path)
    if not nodes or not edges:
        rospy.logerr("unable to load graph data")
        return
    
    # 打印节点和边数据
    info=f"Nodes size :{len(nodes)} ,edges size:{len(edges)}"
    rospy.loginfo(info)

    # 创建发布者
    pub = rospy.Publisher(topic_name, MarkerArray, queue_size=1)
    marker_array = create_marker_array(nodes, edges, frame_id)

    # 发布频率
    rate = rospy.Rate(rate_val)
    while not rospy.is_shutdown():
        pub.publish(marker_array)
        rate.sleep()

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass